// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: donggu

use failure::{err_msg, Error};
use flate2::bufread::GzDecoder;
use std::convert::TryInto;
use std::io::prelude::*;

use pyo3::prelude::*;
use pyo3::types::PyBytes;
use pyo3::PyObjectProtocol;

#[pymodule]
fn topic_record(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<TopicRecordReader>()?;
    Ok(())
}

#[pyclass]
struct TopicRecord {
    #[pyo3(get)]
    queue_seq: u64,
    #[pyo3(get)]
    topic_id: u64,
    #[pyo3(get)]
    topic_seq: u64,
    #[pyo3(get, set)]
    topic_str: String,
    #[pyo3(get)]
    timestamp: i64,
    data: PyObject,
    buf: PyObject,
}

#[pymethods]
impl TopicRecord {
    #[getter]
    fn data(&self, py: Python) -> PyObject {
        self.data.clone_ref(py)
    }
}

#[pyproto]
impl PyObjectProtocol for TopicRecord {
    fn __repr__(&self) -> PyResult<String> {
        Ok(format!(
            "TopicRecord(queue_seq={},topic_id={},topic_seq={},timestamp={},data=<..>)",
            self.queue_seq, self.topic_id, self.topic_seq, self.timestamp,
        ))
    }
}

#[pyclass]
struct TopicRecordReader {
    file: Box<dyn Read>,
    record: Option<TopicRecord>,
    buf: Vec<u8>,
    overwrite_topic_id: Option<u64>,
    topic_str: String,
}

impl TopicRecordReader {
    fn read_impl(&mut self, py: Python) -> Option<TopicRecord> {
        self.read_and_verify_topic_record(py).unwrap()
    }
    fn read_and_verify_magic(&mut self) -> Result<(), Error> {
        self.buf.resize(8, 0);
        const MAGIC: &[u8] = b"\x00\x31\x41\x59\x00\x01\x00\x00";
        self.file.read_exact(&mut self.buf[..MAGIC.len()])?;
        if MAGIC != &self.buf[0..MAGIC.len()] {
            return Err(err_msg("magic dismatch"));
        }
        Ok(())
    }

    fn read_and_verify_topic_record(&mut self, py: Python) -> Result<Option<TopicRecord>, Error> {
        self.buf.resize(44, 0);
        match self.file.read(&mut self.buf[..1]) {
            Ok(0) => return Ok(None),
            Err(err) => Err(err)?,
            _ => {}
        }
        if self.file.read_exact(&mut self.buf[1..44]).is_err() {
            return Err(err_msg("record is shorter than 44 bytes"));
        }

        // 1. length 8 bytes
        let length = u64::from_le_bytes((&self.buf[..8]).try_into()?) as usize;
        if length < 32 {
            return Err(err_msg("record length is smaller than 32"));
        }
        // 2. crc 4 bytes
        let crc_of_length = u32::from_le_bytes((&self.buf[8..12]).try_into()?);
        if masked_crc(&self.buf[..8]) != crc_of_length {
            return Err(err_msg("crc_of_length dismatch"));
        }

        let queue_seq = u64::from_le_bytes((&self.buf[12..20]).try_into()?);
        let topic_id = u64::from_le_bytes((&self.buf[20..28]).try_into()?);
        let topic_seq = u64::from_le_bytes((&self.buf[28..36]).try_into()?);
        let timestamp = i64::from_le_bytes((&self.buf[36..44]).try_into()?);

        self.buf.resize(length + 16, 0);

        self.file
            .read_exact(&mut self.buf[44..44 + (length - 32)])?;
        self.file
            .read_exact(&mut self.buf[length + 12..length + 16])?;
        let crc_of_record = u32::from_le_bytes((&self.buf[length + 12..length + 16]).try_into()?);

        if masked_crc(&self.buf[12..44 + length - 32]) != crc_of_record {
            return Err(err_msg("crc_of_record dismatch"));
        }

        let buf = PyBytes::new(py, &self.buf[44..44 + (length - 32)]).to_object(py);
        let memoryview = unsafe {
            let bytes: &PyBytes = buf.extract(py).unwrap();
            let slice = bytes.as_bytes();
            let memoryview = pyo3::ffi::PyMemoryView_FromMemory(
                std::mem::transmute::<*const u8, *mut i8>(slice.as_ptr()),
                slice.len() as isize,
                pyo3::ffi::PyBUF_READ,
            );
            PyObject::from_owned_ptr(py, memoryview)
        };

        return Ok(Some(TopicRecord {
            queue_seq,
            topic_id: self.overwrite_topic_id.unwrap_or(topic_id),
            topic_seq,
            topic_str: self.topic_str.clone(),
            timestamp,
            data: memoryview,
            buf,
        }));
    }
}

#[pymethods]
impl TopicRecordReader {
    #[new]
    fn new_python(
        obj: &PyRawObject,
        _py: Python,
        path: &str,
        overwrite_topic_id: Option<u64>,
        use_topic_str: Option<bool>,
    ) -> PyResult<()> {
        let file: Box<dyn Read> = if path.ends_with(".gz") {
            Box::new(GzDecoder::new(std::io::BufReader::new(
                std::fs::File::open(path)?,
            )))
        } else {
            Box::new(std::io::BufReader::new(std::fs::File::open(path)?))
        };

        let topic_str = if use_topic_str == Some(true) {
            let path = std::path::Path::new(path);
            let filename = path.file_name().expect("no filename").to_str().expect("not utf-8");
            (&filename[0..filename.find('.').expect("no dot in path")]).to_string()
        } else {
            "".to_string()
        };
        let mut reader = TopicRecordReader {
            file,
            record: None,
            buf: Vec::new(),
            overwrite_topic_id,
            topic_str,
        };
        reader.read_and_verify_magic().expect("wrong magic number");

        obj.init(reader);
        Ok(())
    }

    // peek copies data
    fn peek(&mut self, py: Python) -> PyResult<Option<TopicRecord>> {
        if self.record.is_none() {
            self.record = self.read_impl(py);
        }
        if let Some(record) = &self.record {
            Ok(Some(TopicRecord {
                queue_seq: record.queue_seq,
                topic_id: record.topic_id,
                topic_seq: record.topic_seq,
                topic_str: record.topic_str.clone(),
                timestamp: record.timestamp,
                data: record.data.clone_ref(py),
                buf: record.buf.clone_ref(py),
            }))
        } else {
            Ok(None)
        }
    }

    fn read(&mut self, py: Python) -> Option<TopicRecord> {
        if self.record.is_some() {
            self.record.take()
        } else {
            self.read_impl(py)
        }
    }

    /// there is no code using this function
    fn read_records(&mut self) {
        unreachable!();
    }
}

fn masked_crc(b: &[u8]) -> u32 {
    let value: u32 = crc::crc32::checksum_ieee(b);
    return ((value >> 15) | (value << 17))
        .overflowing_add(0xa282ead8)
        .0
        & 0xffffffff;
}
